import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
plt.rcParams['axes.unicode_minus'] = False    # 正常显示负号

# 实验二：复数方程求解 (简化版)
print("实验二：复数方程求解")
print("="*40)

# 解二次方程
def solve_quadratic(alpha, beta, gamma, delta):
    a, b, c = 1, alpha + 1j*beta, gamma + 1j*delta
    discriminant = b**2 - 4*a*c
    sqrt_disc = np.sqrt(discriminant)
    x1 = (-b + sqrt_disc) / (2*a)
    x2 = (-b - sqrt_disc) / (2*a)
    return x1, x2

print("1. 解二次方程:")
alpha, beta, gamma, delta = 2, -1, 3, 1
x1, x2 = solve_quadratic(alpha, beta, gamma, delta)
print(f"  方程 x² + ({alpha}+{beta}i)x + ({gamma}+{delta}i) = 0")
print(f"  x1 = {x1:.4f}, x2 = {x2:.4f}")

# 验证解
eq_val1 = x1**2 + (alpha + 1j*beta)*x1 + (gamma + 1j*delta)
eq_val2 = x2**2 + (alpha + 1j*beta)*x2 + (gamma + 1j*delta)
print(f"  验证: f(x1) = {eq_val1:.6f}, f(x2) = {eq_val2:.6f}")

# 等边三角形条件
def verify_equilateral(a1, a2, a3):
    left = a1**2 + a2**2 + a3**2
    right = a1*a2 + a2*a3 + a3*a1
    return abs(left - right) < 1e-10

print("\n2. 验证等边三角形条件:")
center = 1 + 1j
radius = 1
a1 = center + radius * np.exp(1j * 0)
a2 = center + radius * np.exp(1j * 2*np.pi/3)
a3 = center + radius * np.exp(1j * 4*np.pi/3)
is_eq = verify_equilateral(a1, a2, a3)
print(f"  三点构成等边三角形: {is_eq}")

# 三角函数倍角公式
print("\n3. 三角函数倍角公式:")
theta = np.pi/6
cos_3theta = np.cos(3*theta)
cos_formula = 4*np.cos(theta)**3 - 3*np.cos(theta)
print(f"  cos(3θ) = {cos_3theta:.6f}")
print(f"  4cos³(θ) - 3cos(θ) = {cos_formula:.6f}")
print(f"  相等: {abs(cos_3theta - cos_formula) < 1e-10}")

# 单位根
def unit_roots(n):
    return [np.exp(2j * np.pi * k / n) for k in range(n)]

def verify_root_sum(n, k):
    omega = np.exp(2j * np.pi / n)
    return sum([omega**(j*k) for j in range(n)])

print("\n4. 单位根:")
roots5 = unit_roots(5)
for k, root in enumerate(roots5):
    print(f"  ω{k} = {root:.4f}")

print("\n  验证根和性质:")
root_sum1 = verify_root_sum(5, 2)  # k不是5的倍数
root_sum2 = verify_root_sum(5, 5)  # k是5的倍数
print(f"  k=2 (非倍数): 1+ω+...+ω⁴ = {root_sum1:.6f}")
print(f"  k=5 (倍数): 1+1+...+1 = {root_sum2:.6f}")

# 可视化单位根
print("\n5. 可视化单位根:")
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))

# 5次单位根
n5 = 5
roots5 = unit_roots(n5)
ax1.plot(np.real(roots5 + [roots5[0]]), np.imag(roots5 + [roots5[0]]), 'b-o', markersize=4)
unit_circle = np.exp(1j * np.linspace(0, 2*np.pi, 100))
ax1.plot(np.real(unit_circle), np.imag(unit_circle), 'k--', alpha=0.5)
ax1.set_xlim(-1.2, 1.2)
ax1.set_ylim(-1.2, 1.2)
ax1.set_title(f'{n5}次单位根')
ax1.grid(True)

# 10次单位根
n10 = 10
roots10 = unit_roots(n10)
ax2.plot(np.real(roots10), np.imag(roots10), 'r-o', markersize=3)
ax2.plot(np.real(unit_circle), np.imag(unit_circle), 'k--', alpha=0.5)
ax2.set_xlim(-1.2, 1.2)
ax2.set_ylim(-1.2, 1.2)
ax2.set_title(f'{n10}次单位根')
ax2.grid(True)

plt.tight_layout()
plt.show()